package com.star.langchain.learn.controller;

import java.util.List;

import com.star.langchain.learn.aiservice.SimpleAiService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.loader.UrlDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.parser.apache.poi.ApachePoiDocumentParser;
import dev.langchain4j.data.document.parser.markdown.MarkdownDocumentParser;
import dev.langchain4j.data.document.parser.yaml.YamlDocumentParser;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;

import lombok.extern.slf4j.Slf4j;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.IngestionResult;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import static dev.langchain4j.data.document.loader.ClassPathDocumentLoader.loadDocument;

@RestController
@RequestMapping("/rag")
@Slf4j
public class RAGController {

    @Autowired
    private EmbeddingModel embeddingModel;


    @Autowired
    private EmbeddingStoreIngestor embeddingStoreIngestor;

    @GetMapping("/fromTxt")
    public String parseDocument(String msg, Long memoryId) {

        String documentPath = "knowledges/document.txt";
        // Now, let's load a document that we want to use for RAG.
        // We will use the terms of use from an imaginary car rental company, "Miles of Smiles".
        // For this example, we'll import only a single document, but you can load as many as you need.
        // LangChain4j offers built-in support for loading documents from various sources:
        // File System, URL, Amazon S3, Azure Blob Storage, GitHub, Tencent COS.
        // Additionally, LangChain4j supports parsing multiple document types:
        // text, pdf, doc, xls, ppt.
        // However, you can also manually import your data from other sources.
        DocumentParser documentParser = new TextDocumentParser();
        Document document = loadDocument(documentPath, documentParser);

        // Now, we need to split this document into smaller segments, also known as "chunks."
        // This approach allows us to send only relevant segments to the LLM in response to a user query,
        // rather than the entire document. For instance, if a user asks about cancellation policies,
        // we will identify and send only those segments related to cancellation.
        // A good starting point is to use a recursive document splitter that initially attempts
        // to split by paragraphs. If a paragraph is too large to fit into a single segment,
        // the splitter will recursively divide it by newlines, then by sentences, and finally by words,
        // if necessary, to ensure each piece of text fits into a single segment.
        DocumentSplitter splitter = DocumentSplitters.recursive(300, 0);
        List<TextSegment> segments = splitter.split(document);


        // Now, we need to embed (also known as "vectorize") these segments.
        // Embedding is needed for performing similarity searches.
        // For this example, we'll use a local in-process embedding model, but you can choose any supported model.
        // Langchain4j currently supports more than 10 popular embedding model providers.
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();


        // Next, we will store these embeddings in an embedding store (also known as a "vector database").
        // This store will be used to search for relevant segments during each interaction with the LLM.
        // For simplicity, this example uses an in-memory embedding store, but you can choose from any supported store.
        // Langchain4j currently supports more than 15 popular embedding stores.


        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings, segments);

        // We could also use EmbeddingStoreIngestor to hide manual steps above behind a simpler API.
        // See an example of using EmbeddingStoreIngestor in _01_Advanced_RAG_with_Query_Compression_Example.

        IngestionResult ingest = embeddingStoreIngestor.ingest(document);

        // The content retriever is responsible for retrieving relevant content based on a user query.
        // Currently, it is capable of retrieving text segments, but future enhancements will include support for
        // additional modalities like images, audio, and more.
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(2) // on each interaction we will retrieve the 2 most relevant segments
                .minScore(0.5) // we want to retrieve segments at least somewhat similar to user query
                .build();




//        // The final step is to build our AI Service,
//        // configuring it to use the components we've created above.
//        SimpleAiService ragAiService = AiServices.builder(SimpleAiService.class)
//                .chatModel(chatModel)
//                .contentRetriever(contentRetriever)
//                .chatMemoryProvider(chatMemoryProvider)
//                .build();
//
//        String chat = ragAiService.chat(msg, memoryId);
//
//        return chat;
        return "success";
    }

    /**
     * 读系统内部文档存入向量数据库
     */
    @GetMapping("fromPdf")
    public String loadFromPdf() {
        Document document = FileSystemDocumentLoader.loadDocument("/Users/star/IdeaProjects/learn/src/main/resources/knowledges/highway.pdf", new ApachePdfBoxDocumentParser());

        IngestionResult ingest = embeddingStoreIngestor.ingest(document);
        return "success";

    }
    /**
     * 读取外部文档
     * 根据文档格式选择不同的解析器，如果解析失败则使用通用解析器
     */
    @GetMapping("fromUrl")
    public String loadFromUrl(String url) {
        try {
            // 根据URL后缀判断文档类型并选择合适的解析器
            DocumentParser parser = selectParserByUrl(url);
            log.info("使用解析器: {} 加载URL: {}", parser.getClass().getSimpleName(), url);
            
            // 尝试使用选定的解析器加载文档
            Document document = UrlDocumentLoader.load(url, parser);
            
            // 将文档存入向量数据库
            IngestionResult ingest = embeddingStoreIngestor.ingest(document);
            log.info("文档加载成功，已存入向量数据库，URL: {}", url);
            
            return "success";
        } catch (Exception e) {
            // 如果特定解析器失败，尝试使用通用解析器
            log.warn("使用特定解析器失败，尝试使用通用解析器，URL: {}, 错误: {}", url, e.getMessage());
            try {
                DocumentParser fallbackParser = new ApacheTikaDocumentParser();
                Document document = UrlDocumentLoader.load(url, fallbackParser);
                IngestionResult ingest = embeddingStoreIngestor.ingest(document);
                log.info("使用通用解析器成功，文档已存入向量数据库，URL: {}", url);
                return "success";
            } catch (Exception fallbackException) {
                log.error("使用通用解析器也失败，URL: {}", url, fallbackException);
                throw new RuntimeException("文档加载失败: " + fallbackException.getMessage(), fallbackException);
            }
        }
    }

    /**
     * 根据URL的文件扩展名选择合适的解析器
     *
     * @param url 文档URL
     * @return 合适的文档解析器
     */
    private DocumentParser selectParserByUrl(String url) {
        if (url == null || url.isEmpty()) {
            // 如果无法判断，默认使用通用解析器
            return new ApacheTikaDocumentParser();
        }

        // 提取文件扩展名（转换为小写以便比较）
        String lowerUrl = url.toLowerCase();
        
        // PDF文件
        if (lowerUrl.endsWith(".pdf")) {
            return new ApachePdfBoxDocumentParser();
        }
        
        // Office文档 (Word, Excel, PowerPoint)
        if (lowerUrl.endsWith(".doc") || lowerUrl.endsWith(".docx") ||
            lowerUrl.endsWith(".xls") || lowerUrl.endsWith(".xlsx") ||
            lowerUrl.endsWith(".ppt") || lowerUrl.endsWith(".pptx")) {
            return new ApachePoiDocumentParser();
        }
        
        // Markdown文件
        if (lowerUrl.endsWith(".md") || lowerUrl.endsWith(".markdown")) {
            return new MarkdownDocumentParser();
        }
        
        // YAML文件
        if (lowerUrl.endsWith(".yaml") || lowerUrl.endsWith(".yml")) {
            return new YamlDocumentParser();
        }
        
        // 纯文本文件 (TXT, HTML等)
        if (lowerUrl.endsWith(".txt") || lowerUrl.endsWith(".html") || 
            lowerUrl.endsWith(".htm") || lowerUrl.endsWith(".xml")) {
            return new TextDocumentParser();
        }
        
        // 无法识别格式，使用通用解析器
        log.debug("无法根据URL后缀识别文档类型，将使用通用解析器: {}", url);
        return new ApacheTikaDocumentParser();
    }


}
